import DRR.drr_generator as drr
import numpy as np
import matplotlib.pyplot as plt
import time
import random
from PIL import Image
import scipy.stats as stats
import json
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from utils.image_process import caculate_NCC, get_new_drr_numpy
from tqdm import tqdm
from collections import OrderedDict

plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置字体
plt.rcParams["axes.unicode_minus"] = False


# 生成训练数据
def get_data(nums=50000, mode=None, saving_path=None):
    """

    :param nums: 样本数
    :param mode:采样模式[正位或侧位，样本分布]
    :param saving_path:文件保存地址
    :return:
    """
    data_generator = drr.Projector(directory="C:/Users/adminTKJ/Desktop/RLIR_sumup/CT_data")
    t1 = time.time()
    # 设定采样范围
    if mode[0] == '标准正位':
        range_rx = [0, 5]
        range_ry = [270, 10]
        range_rz = [90, 5]
        range_tx = 25
        range_ty = 50
        range_tz = 25
    elif mode[0] == '标准侧位':
        range_rx = [0, 5]
        range_ry = [90, 5]
        range_rz = [180, 5]
        range_tx = 50
        range_ty = 25
        range_tz = 25
    else:
        return print("模式不正确，请重新输入")
    # 字典文件保存文件标签
    im_dict = {}
    for i in range(nums):
        if mode[1] == 'mean':
            # 均匀采样
            rx_noise = np.random.uniform(-1, 1) * range_rx[1]
            ry_noise = np.random.uniform(-1, 1) * range_ry[1]
            rz_noise = np.random.uniform(-1, 1) * range_rz[1]
            tx_noise = np.random.uniform(-1, 1) * range_tx
            ty_noise = np.random.uniform(-1, 1) * range_ty
            tz_noise = np.random.uniform(-1, 1) * range_tz
        elif mode[1] == 'gauss':
            mu, sigma = 0, 0.5
            lower, upper = mu - 2 * sigma, mu + 2 * sigma  # 截断在[μ-2σ, μ+2σ]
            X = stats.truncnorm((lower - mu) / sigma, (upper - mu) / sigma, loc=mu, scale=sigma)
            rx_noise = X.rvs(1)[0] * range_rx[1]
            ry_noise = X.rvs(1)[0] * range_ry[1]
            rz_noise = X.rvs(1)[0] * range_rz[1]
            tx_noise = X.rvs(1) * range_tx
            ty_noise = X.rvs(1) * range_ty
            tz_noise = X.rvs(1) * range_tz
        else:
            return print("采样分布不正确，请重新输入")
        angle_noise = np.array([rx_noise, ry_noise, rz_noise], dtype=np.float32)
        trans_noise = np.array([tx_noise, ty_noise, tz_noise], dtype=np.float32)
        data_generator.set_rotation(range_rx[0], range_ry[0], range_rz[0])
        im = data_generator.project(angle_noise, trans_noise)
        # 归一±1之间
        rx, ry, rz = rx_noise / range_rx[1], ry_noise / range_ry[1], rz_noise / range_rz[1]
        tx, ty, tz = tx_noise/range_tx, ty_noise/range_ty, tz_noise/range_tz,
        t_info = str(round(tx, 2)) + "_" + str(round(ty, 2)) + "_" + str(round(tz, 2))
        im_info = "x{0}_y{1}_z{2}_{3}.npy".format(round(rx, 2), round(ry, 2), round(rz, 2), t_info)
        im_dict[im_info] = {
            "name": im_info,
            "angle_x": rx,
            "angle_y": ry,
            "angle_z": rz,
            "tx": tx,
            "ty": ty,
            "tz": tz,
        }
        np.save(saving_path + "/DRR/{0}".format(im_info), im)
        print("正在生成第{0}组图片，文件信息{1}，当前时间:{2}".format(i, im_dict[im_info],
                                                                    time.asctime(time.localtime(time.time()))))
    # 将文件写入json
    with open(saving_path + "/label.json", "w") as f:
        json.dump(im_dict, f)
    print("加载入文件完成...")
    t2 = time.time()
    print("用时：{0}".format(t2 - t1))


if __name__ == "__main__":
    # path = 'C:/Users/adminTKJ/Desktop/RLIR_sumup/data/CT/截断高斯采样
    # /train/标准正位/DRR/x0.38_y0.72_z-0.48_-0.43_-0.2_-0.13.npy.npy'
    # im = np.load(path)
    # data_generator = Projector()
    # data_generator.load_ct_images(directory="C:/Users/adminTKJ/Desktop/CT_投影/CT_data")
    # 标准正位 0 270 90（本次实验用到的）或者0 90 270
    # 标准侧位 0 270 0或者0 90 180（本次实验用到的）
    # data_generator.set_rotation(0, 270, 90)
    # data_generator.project(0, 0, 0, mode='display', save=False)
    get_data(nums=50000, mode=['标准正位', 'mean'],
             saving_path='C:/Users/adminTKJ/Desktop/RLIR_sumup/data/CT/均匀采样/train/标准正位')
